"""
Token Authentication Middleware for WebSocket
支持从查询参数中获取Token进行WebSocket认证
"""
from channels.db import database_sync_to_async
from channels.middleware import BaseMiddleware
from django.contrib.auth.models import AnonymousUser
from rest_framework.authtoken.models import Token
from urllib.parse import parse_qs


@database_sync_to_async
def get_user_from_token(token_key):
    """根据token获取用户"""
    try:
        token = Token.objects.get(key=token_key)
        return token.user
    except Token.DoesNotExist:
        return AnonymousUser()


class TokenAuthMiddleware(BaseMiddleware):
    """
    Token认证中间件
    从URL查询参数 ?token=xxx 中获取token并认证用户
    """
    
    async def __call__(self, scope, receive, send):
        # 解析查询参数
        query_string = scope.get('query_string', b'').decode()
        query_params = parse_qs(query_string)
        
        # 获取token
        token_key = query_params.get('token', [None])[0]
        
        if token_key:
            # 从token获取用户
            scope['user'] = await get_user_from_token(token_key)
        else:
            # 没有token，设置为匿名用户
            scope['user'] = AnonymousUser()
        
        return await super().__call__(scope, receive, send)
